/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "parcsr_matrix.h"
#include "parcsr_matvec.h"
#include "parcsr_matmultivec.h"

namespace YHAMG
{

ParCSR_Matrix::ParCSR_Matrix()
	: num_neighbors(0),
	neighbor_ranks(0),
	Ext_ptr(0),
	Ext_map(0),
	send_ptr(0),
	send_list(0),
	send_buffer(0)
{
}

ParCSR_Matrix::ParCSR_Matrix(MPI_Comm _comm, int LocLoc_rows, int LocLoc_cols, int LocExt_cols, 
	int* LocLoc_rowptr, int* LocLoc_colind, double* LocLoc_values, int LocLoc_ref, 
	int* LocExt_rowptr, int* LocExt_colind, double* LocExt_values, int LocExt_ref,
	int _num_neighbors, int* _neighbor_ranks, int* _Ext_ptr, int* _Ext_map)
	: Par_Operator(_comm),
	LocLoc(LocLoc_rows, LocLoc_cols, LocLoc_rowptr, LocLoc_colind, LocLoc_values, LocLoc_ref),
	LocExt(LocLoc_rows, LocExt_cols, LocExt_rowptr, LocExt_colind, LocExt_values, LocExt_ref),
	num_neighbors(_num_neighbors),
	neighbor_ranks(_neighbor_ranks),
	Ext_ptr(_Ext_ptr),
	Ext_map(_Ext_map),
	send_ptr(0),
	send_list(0),
	send_buffer(0)
{
}

ParCSR_Matrix::ParCSR_Matrix(const ParCSR_Matrix& A)
	: Par_Operator(A.comm),
	LocLoc(A.LocLoc),
	LocExt(A.LocExt),
	num_neighbors(A.num_neighbors),
	neighbor_ranks(new int[A.num_neighbors]),
	Ext_ptr(new int[num_neighbors + 1]),
	Ext_map(new int[LocExt.size[1]]),
	send_ptr(0),
	send_list(0),
	send_buffer(0)
{
	for (int r = 0; r < num_neighbors; ++r)
		neighbor_ranks[r] = A.neighbor_ranks[r];
	for (int r = 0; r <= num_neighbors; ++r)
		Ext_ptr[r] = A.Ext_ptr[r];
	for (int i = 0; i < LocExt.size[1]; ++i)
		Ext_map[i] = A.Ext_map[i];
}

ParCSR_Matrix::ParCSR_Matrix(ParCSR_Matrix&& A)
	: Par_Operator(A.comm),
	LocLoc(A.LocLoc.size[0], A.LocLoc.size[1], A.LocLoc.rowptr, A.LocLoc.colind, A.LocLoc.values, A.LocLoc.ref),
	LocExt(A.LocExt.size[0], A.LocExt.size[1], A.LocExt.rowptr, A.LocExt.colind, A.LocExt.values, A.LocExt.ref),
	num_neighbors(A.num_neighbors),
	neighbor_ranks(A.neighbor_ranks),
	Ext_ptr(A.Ext_ptr),
	Ext_map(A.Ext_map),
	send_ptr(A.send_ptr),
	send_list(A.send_list),
	send_buffer(A.send_buffer),
	x_Ext(A.x_Ext.size, A.x_Ext.values, A.x_Ext.ref)
{
	A.LocLoc.ref = 1;
	A.LocExt.ref = 1;
	A.neighbor_ranks = 0;
	A.Ext_ptr = 0;
	A.Ext_map = 0;
	A.send_ptr = 0;
	A.send_list = 0;
	A.send_buffer = 0;
	A.x_Ext.ref = 1;
}

ParCSR_Matrix::~ParCSR_Matrix()
{
	if (neighbor_ranks) delete[] neighbor_ranks;
	if (Ext_ptr) delete[] Ext_ptr;
	if (Ext_map) delete[] Ext_map;
	if (send_ptr) delete[] send_ptr;
	if (send_list) delete[] send_list;
	if (send_buffer) delete[] send_buffer;
}

ParCSR_Matrix& ParCSR_Matrix::operator=(const ParCSR_Matrix& A)
{
	Free();

	comm = A.comm;
	LocLoc = A.LocLoc;
	LocExt = A.LocExt;
	num_neighbors = A.num_neighbors;

	neighbor_ranks = new int[num_neighbors];
	Ext_ptr = new int[num_neighbors + 1];
	Ext_map = new int[LocExt.size[1]];

	for (int r = 0; r < num_neighbors; ++r)
		neighbor_ranks[r] = A.neighbor_ranks[r];
	for (int r = 0; r <= num_neighbors; ++r)
		Ext_ptr[r] = A.Ext_ptr[r];
	for (int i = 0; i < LocExt.size[1]; ++i)
		Ext_map[i] = A.Ext_map[i];
	
	return *this;
}

ParCSR_Matrix& ParCSR_Matrix::operator=(ParCSR_Matrix&& A)
{
	Free();

	comm = A.comm;
	LocLoc.ref = A.LocLoc.ref;
	LocLoc.size[0] = A.LocLoc.size[0];
	LocLoc.size[1] = A.LocLoc.size[1];
	LocLoc.rowptr = A.LocLoc.rowptr;
	LocLoc.colind = A.LocLoc.colind;
	LocLoc.values = A.LocLoc.values;
	LocExt.ref = A.LocExt.ref;
	LocExt.size[0] = A.LocExt.size[0];
	LocExt.size[1] = A.LocExt.size[1];
	LocExt.rowptr = A.LocExt.rowptr;
	LocExt.colind = A.LocExt.colind;
	LocExt.values = A.LocExt.values;
	num_neighbors = A.num_neighbors;
	neighbor_ranks = A.neighbor_ranks;
	Ext_ptr = A.Ext_ptr;
	Ext_map = A.Ext_map;
	send_ptr = A.send_ptr;
	send_list = A.send_list;
	send_buffer = A.send_buffer;
	x_Ext.ref = A.x_Ext.ref;
	x_Ext.size = A.x_Ext.size;
	x_Ext.values = A.x_Ext.values;

	A.LocLoc.ref = 1;
	A.LocExt.ref = 1;
	A.neighbor_ranks = 0;
	A.Ext_ptr = 0;
	A.Ext_map = 0;
	A.send_ptr = 0;
	A.send_list = 0;
	A.send_buffer = 0;
	A.x_Ext.ref = 1;

	return *this;
}

void ParCSR_Matrix::Free()
{
	if (neighbor_ranks) delete[] neighbor_ranks;
	if (Ext_ptr) delete[] Ext_ptr;
	if (Ext_map) delete[] Ext_map;
	if (send_ptr) delete[] send_ptr;
	if (send_list) delete[] send_list;
	if (send_buffer) delete[] send_buffer;

	LocLoc.Free();
	LocExt.Free();
	num_neighbors = 0;
	neighbor_ranks = 0;
	Ext_ptr = 0;
	Ext_map = 0;
	send_ptr = 0;
	send_list = 0;
	send_buffer = 0;
	x_Ext.Free();
}

void ParCSR_Matrix::Reference(const ParCSR_Matrix& A)
{
	Free();

	comm = A.comm;
	LocLoc.Reference(A.LocLoc);
	LocExt.Reference(A.LocExt);
	num_neighbors = A.num_neighbors;

	neighbor_ranks = new int[num_neighbors];
	Ext_ptr = new int[num_neighbors + 1];
	Ext_map = new int[LocExt.size[1]];

	for (int r = 0; r < num_neighbors; ++r)
		neighbor_ranks[r] = A.neighbor_ranks[r];
	for (int r = 0; r <= num_neighbors; ++r)
		Ext_ptr[r] = A.Ext_ptr[r];
	for (int i = 0; i < LocExt.size[1]; ++i)
		Ext_map[i] = A.Ext_map[i];
}

#define MPI_TAG 100

void ParCSR_Matrix::ExchangeHalo(const Vector& x_Loc, const Vector& x_Ext) const
{
	double* xv_Loc = x_Loc.values;
	double* xv_Ext = x_Ext.values;

	MPI_Request* recv_requests = new MPI_Request[num_neighbors];
	MPI_Request* send_requests = new MPI_Request[num_neighbors];
	MPI_Status status;

	for (int r = 0; r < num_neighbors; ++r)
		if (Ext_ptr[r + 1] > Ext_ptr[r])
			MPI_Irecv(xv_Ext + Ext_ptr[r], Ext_ptr[r + 1] - Ext_ptr[r], MPI_DOUBLE, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			for (int i = send_ptr[r]; i < send_ptr[r + 1]; ++i)
				send_buffer[i] = xv_Loc[send_list[i]];
			MPI_Isend(send_buffer + send_ptr[r], send_ptr[r + 1] - send_ptr[r], MPI_DOUBLE, neighbor_ranks[r], MPI_TAG, comm, send_requests + r);
		}
	}

	for (int r = 0; r < num_neighbors; ++r)
		if (Ext_ptr[r + 1] > Ext_ptr[r])
			MPI_Wait(recv_requests + r, &status);

	for (int r = 0; r < num_neighbors; ++r)
		if (send_ptr[r + 1] > send_ptr[r])
			MPI_Wait(send_requests + r, &status);

	delete[] recv_requests;
	delete[] send_requests;
}

void ParCSR_Matrix::SetupHalo()
{
	if (send_ptr) delete[] send_ptr;
	if (send_list) delete[] send_list;
	if (send_buffer) delete[] send_buffer;

	x_Ext.Resize(LocExt.size[1]);

	MPI_Request* recv_requests = new MPI_Request[num_neighbors];
	MPI_Request* send_requests = new MPI_Request[num_neighbors];
	MPI_Status status;

	send_ptr = new int[num_neighbors + 1];

	send_ptr[0] = 0;
	for (int r = 0; r < num_neighbors; ++r)
		MPI_Irecv(send_ptr + r + 1, 1, MPI_INT, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

	for (int r = 0; r < num_neighbors; ++r)
	{
		int cnt = Ext_ptr[r + 1] - Ext_ptr[r];
		MPI_Send(&cnt, 1, MPI_INT, neighbor_ranks[r], MPI_TAG, comm);
	}

	for (int r = 0; r < num_neighbors; ++r)
		MPI_Wait(recv_requests + r, &status);

	for (int r = 0; r < num_neighbors; ++r)
		send_ptr[r + 1] += send_ptr[r];

	int _num_neighbors = 0;

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r] || Ext_ptr[r + 1] > Ext_ptr[r])
		{
			neighbor_ranks[_num_neighbors++] = neighbor_ranks[r];
			send_ptr[_num_neighbors] = send_ptr[r + 1];
			Ext_ptr[_num_neighbors] = Ext_ptr[r + 1];
		}
	}

	num_neighbors = _num_neighbors;

	send_list = new int[send_ptr[num_neighbors]];
	send_buffer = new double[send_ptr[num_neighbors]];

	for (int r = 0; r < num_neighbors; ++r)
		if (send_ptr[r + 1] > send_ptr[r])
			MPI_Irecv(send_list + send_ptr[r], send_ptr[r + 1] - send_ptr[r], MPI_INT, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

	for (int r = 0; r < num_neighbors; ++r)
		if (Ext_ptr[r + 1] > Ext_ptr[r])
			MPI_Isend(Ext_map + Ext_ptr[r], Ext_ptr[r + 1] - Ext_ptr[r], MPI_INT, neighbor_ranks[r], MPI_TAG, comm, send_requests + r);

	for (int r = 0; r < num_neighbors; ++r)
		if (send_ptr[r + 1] > send_ptr[r])
			MPI_Wait(recv_requests + r, &status);

	for (int r = 0; r < num_neighbors; ++r)
		if (Ext_ptr[r + 1] > Ext_ptr[r])
			MPI_Wait(send_requests + r, &status);

	delete[] recv_requests;
	delete[] send_requests;
}

int ParCSR_Matrix::InSize() const
{
	return LocLoc.size[1];
}

int ParCSR_Matrix::OutSize() const
{
	return LocLoc.size[0];
}

void ParCSR_Matrix::Apply(const Par_Vector& x, const Par_Vector& y) const
{
	ParCSR_MatVec(*this, x, y);
}

void ParCSR_Matrix::Apply(const Par_MultiVector& X, const Par_MultiVector& Y) const
{
	ParCSR_MatMultiVec(*this, X, Y);
}

}